/**
 * 2018年12月16日上午6:34:15
 * @author 谢进财
 */
package com.baihang.project.ds.webSocket.handler;

import java.io.IOException;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;

import org.apache.commons.lang3.concurrent.BasicThreadFactory;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketHandler;
import org.springframework.web.socket.WebSocketMessage;
import org.springframework.web.socket.WebSocketSession;

/**
 * @author whati
 *
 */

public class MyWebSocketHandler implements WebSocketHandler {
	    private static final Logger logger = LoggerFactory.getLogger(MyWebSocketHandler.class);
	    private static Map<String, WebSocketSession> userSocketSessionMap = new ConcurrentHashMap<>();

	    /**
	     * 建立连接后
	     */
	    @Override
	    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
	        String userName = (String) session.getAttributes().get("userName");
	        logger.info("Session {} connected.", userName);

	        // 如果是新 Session，记录进 Map
	        boolean isNewUser = true;
	        for (Entry<String, WebSocketSession> entry: userSocketSessionMap.entrySet()) {
	            String key = entry.getKey();
	            if (key.equals(userName)) {
	                isNewUser = false;
	                break;
	            }
	        }
	        if (isNewUser) {	           
	            userSocketSessionMap.put(userName, session);
	        }
	        logger.info("当前在线用户数: {}", userSocketSessionMap.size());
	    }

	    /**
	     * 消息处理，在客户端通过Websocket API发送的消息会经过这里，然后进行相应的处理
	     */
	    @Override
	    public void handleMessage(WebSocketSession session, WebSocketMessage<?> message) throws Exception {
	        if (message.getPayloadLength() == 0) {
	            return;
	        }
	        new TextMessage("连接成功");
	    }

	    /**
	     * 消息传输错误处理
	     */
	    @Override
	    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
	        if (session.isOpen()) {
	            session.close();
	        }
	        if(userSocketSessionMap.containsValue(session)) {
	        	userSocketSessionMap.remove((String) session.getAttributes().get("userName"));
	        }
	    }

	    /**
	     * 关闭连接后
	     */
	    @Override
	    public void afterConnectionClosed(WebSocketSession session, CloseStatus closeStatus) throws Exception {
	        logger.info("Session {} disconnected. Because of {}", session.getId(), closeStatus);
	        if(userSocketSessionMap.containsValue(session)) {
	        	userSocketSessionMap.remove((String) session.getAttributes().get("userName"));
	        }
	        
	        logger.info("当前在线用户数: {}", userSocketSessionMap.size());
	    }

	    @Override
	    public boolean supportsPartialMessages() {
	        return false;
	    }

	    /**
	     * 给所有在线用户发送消息
	     *
	     * @param message
	     * @throws IOException
	     */
	    public void broadcast(final TextMessage message) throws IOException {
	    	System.err.println("群发");
	        // 多线程群发
	        for(Entry<String, WebSocketSession> entry: userSocketSessionMap.entrySet()) {
	        	WebSocketSession  session = entry.getValue();
	                if (session.isOpen()) {
	                    ScheduledExecutorService executorService = new ScheduledThreadPoolExecutor(1,
	                        new BasicThreadFactory.Builder().namingPattern("socket-schedule-pool-%d").daemon(true).build());
	                  
	                        executorService.execute(new Runnable() {
	                            @Override
	                            public void run() {
	                                try {
	                                    if (session.isOpen()) {
	                                        logger.debug("broadcast output:" + message.toString());
	                                        session.sendMessage(message);
	                                    }
	                                } catch (IOException e) {
	                                    e.printStackTrace();
	                                }
	                            }
	                        });
	                  
	                }
	            
	        }
	    }

	    /**
	     * 给某个用户发送消息
	     *
	     * @param userName
	     * @param message
	     * @throws IOException
	     */
	    private void sendMessageToUser(String userName, TextMessage message) throws IOException {
	    	WebSocketSession session = userSocketSessionMap.get(userName);	        
	        session.sendMessage(message);	                    
	    }

}
